#ifndef __COROUTIEN_THREAD_POOL_H__
#define __COROUTIEN_THREAD_POOL_H__

#include <atomic>
#include <condition_variable>
#include <coroutine>
#include <iostream>
#include <mutex>
#include <optional>
#include <queue>
#include <stdexcept>
#include <thread>

namespace coroutine_thread_pool
{

	template <typename T>
	struct threadsafe_queue_t
	{
		threadsafe_queue_t() {}

		void put(T task)
		{
			std::unique_lock<std::mutex> lk(m_m);
			m_queue.emplace(task);
			m_cv.notify_one();
		}

		std::optional<T> take()
		{
			std::unique_lock<std::mutex> lk(m_m);
			m_cv.wait(lk, [q = this]
				{ return q->m_must_return_nullptr.test() || !q->m_queue.empty(); });

			if (m_must_return_nullptr.test())
				return {};

			T ret = m_queue.front();
			m_queue.pop();

			return ret;
		}

		void destroy()
		{
			m_must_return_nullptr.test_and_set();
			m_cv.notify_all();
		}

	private:
		std::queue<T> m_queue;
		std::mutex m_m;
		std::condition_variable m_cv;

		std::atomic_flag m_must_return_nullptr = false;
	};

	struct thread_pool_t
	{
		struct promise_type;
		struct awaitable
		{
			bool await_ready()
			{
				return false;
			}

			void await_suspend(std::coroutine_handle<> h)
			{
				m_thread_pool->m_queue.put(h);
			}

			void await_resume() {}

			thread_pool_t* m_thread_pool;
		};

		struct wrapper
		{
			struct promise_type
			{
				wrapper get_return_object() { return {}; }
				std::suspend_never initial_suspend() { return {}; }
				std::suspend_never final_suspend() noexcept { return {}; }
				void return_void() {}
				void unhandled_exception() {}

				// 构造函数
				promise_type() {}
			};
		};

		thread_pool_t(int thread_number)
		{
			for (int i = 0; i < thread_number; ++i)
			{
				m_threads.emplace(std::jthread(
					[this]
					{ this->worker(); }));
			}
		}

		template <typename T>
		wrapper submit(T task)
		{
			co_await awaitable(this);

			task();
		}

		void stop()
		{
			m_queue.destroy();
			while (!m_threads.empty())
			{
				m_threads.pop();
			}
		}

		threadsafe_queue_t<std::coroutine_handle<>> m_queue;

	private:
		void worker()
		{
			while (auto task = m_queue.take())
			{
				task.value().resume();
			}
		}

		std::queue<std::jthread> m_threads;
	};

	void test()
	{
		using namespace std::chrono_literals;
		auto tpool = thread_pool_t(3);
		for (int i = 0; i < 7; ++i)
		{
			tpool.submit(
				[i]
				{
					std::this_thread::sleep_for(2s);
					std::cout << "* Task " << i << std::endl;
				});
		}

		std::this_thread::sleep_for(3s);
		tpool.stop();
	}
}

#endif